import argparse
from datetime import datetime, timedelta
import baostock as bs
import pymysql


DB_CONFIG = {
    'user': 'komi',
    'password': 'komi2025',
    'database': 'komi',
    'host': 'localhost',
    'port': 3306
}


class Args:
    
    @staticmethod
    def parse():
        parser = argparse.ArgumentParser()
        yesterday = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')
        parser.add_argument('-d', '--date', type=str, default=yesterday, help='Date in YYYY-MM-DD format')
        parser.add_argument('--start-date', type=str, default=yesterday, help='Date in YYYY-MM-DD format')
        parser.add_argument('--end-date', type=str, default=yesterday, help='Date in YYYY-MM-DD format')
        return parser.parse_args()


class StockApi:

    def __init__(self):
        bs.login()
        self.bs = bs

    def to_null(self, s):
        if not s:
            return None
        else:
            return s

    def __getattr__(self, name):
        attr = getattr(self.bs, name, None)
        if callable(attr):
            # 定义装饰器
            def to_iterable(func):
                def wrapper(*args, **kwargs):
                    fields_map = kwargs.get('fields_map')
                    if fields_map:
                        del kwargs['fields_map'] # 接口本身不需要这个参数
                    
                    rs = func(*args, **kwargs)
                    
                    # 将原始字段映射成mysql表字段
                    # 若无 fields_map 参数则不处理

                    if fields_map:
                        fields = []
                        for field in rs.fields:
                            fields.append(fields_map[field])
                    else:
                        fields = rs.fields
                    while rs.next():
                        row = [self.to_null(col) for col in rs.get_row_data()]
                        yield dict(zip(fields, row))
                return wrapper
            return to_iterable(attr)
        else:
            raise AttributeError("'StockApi' object has no attribute '{name}'")


class Database:

    def execute(self, sql, args, many=False):
        db_config = DB_CONFIG.copy()
        db_config['cursorclass'] = pymysql.cursors.DictCursor
        with pymysql.connect(**db_config) as conn:
            cursor = conn.cursor()
            if many:
                cursor.executemany(sql, args)    
            else:
                cursor.execute(sql, args)
            conn.commit()
            try:
                result = cursor.fetchall()
            except pymysql.err.ProgrammingError:
                result = None
            return result

    def write(self, table_name: str, data: dict):
        """
        作用：将接口数据写入到数据库
        原则：st 层保留接口的所有字段，表字段和接口字段保持一致，不一致则抛出异常
             所有的表都必须有 dt 日期字段，没有例外

        table_name: 目标表
        data: 待插入的数据，用 StockApi.<接口> 即可获得
        """
        fields = ','.join([f'{key}' for key in data[0]])
        placeholders = ','.join([f'%({key})s' for key in data[0]])
        # print(f'insert {table_name} ({fields}) values ({placeholders})')
        self.execute(f'insert {table_name} ({fields}) values ({placeholders})', data, many=True)


if __name__ == '__main__':
    api = StockApi()
    rs = api.query_stock_basic(code="sh.600000")
    api.write()